HyperPod in Studio のダッシュボードを眺めてみた
2024.12.28こんにちは！AWS 事業本部コンサルティング部のたかくに（@takakuni_）です。
Amazon SageMaker Studio で SageMaker HyperPod の情報を表示する HyperPod in Studio を試してみる機会があったためブログにします。
HyperPod in Studio
説明が重複しますが、HyperPod in Studio は Amazon SageMaker HyperPod を SageMaker Studio 上で管理する機能です。
主にクラスターの情報やメトリクスなどを可視化、タスクのジョブを GUI で管理するこことができます。
re:Invent 2024 のタスクガバナンスの機能に合わせて、しれっと紹介されていました。
Additionally, HyperPod is now integrated with Amazon SageMaker Studio, bringing task governance and other HyperPod capabilities into the Studio environment.
タスクガバナンスは EKS オーケストレータの場合のみ利用できますが、 HyperPod in Studio は Slurm, EKS のどちらものオーケーストレータにサポートしています。
やってみる
以下の構成で Pytorch による DDP を実行した場合の HyperPod in Studio の様子を見てみたいと思います。
IAM
HyperPod in Studio を利用するには SageMaker Studio の実行ロールに追加の権限が必要です。
オーケストレータによって権限が異なるため注意です。私の場合はセットアップに Terraform 、オーケストレータは Slurm を利用したため、以下のようなコードになりました。
studio.tf
###################################################
# IAM Role for SageMaker Domain
###################################################
resource "aws_iam_role" "studio" {
name = "${local.prefix}-studio-role"
assume_role_policy = jsonencode({
Version = "2012-10-17",
Statement = [
{
Effect = "Allow",
Principal = {
Service = "sagemaker.amazonaws.com"
},
Action = "sts:AssumeRole"
}
]
})
}
# https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-migrate-ui.html
data "aws_iam_policy_document" "studio_basic" {
statement {
sid = "SMStudioUserProfileAppPermissionsCreateAndDelete"
effect = "Allow"
actions = [
"sagemaker:CreateApp",
"sagemaker:DeleteApp"
]
resources = [
"arn:aws:sagemaker:${local.region}:${local.account_id}:app/*"
]
condition {
test = "Null"
variable = "sagemaker:OwnerUserProfileArn"
values = ["true"]
}
}
statement {
sid = "SMStudioCreatePresignedDomainUrlForUserProfile"
effect = "Allow"
actions = [
"sagemaker:CreatePresignedDomainUrl"
]
resources = [
aws_sagemaker_user_profile.this.arn
]
}
statement {
sid = "SMStudioAppPermissionsListAndDescribe"
effect = "Allow"
actions = [
"sagemaker:ListApps",
"sagemaker:ListDomains",
"sagemaker:ListUserProfiles",
"sagemaker:ListSpaces",
"sagemaker:DescribeApp",
"sagemaker:DescribeDomain",
"sagemaker:DescribeUserProfile",
"sagemaker:DescribeSpace"
]
resources = ["*"]
}
statement {
sid = "SMStudioAppPermissionsTagOnCreate"
effect = "Allow"
actions = [
"sagemaker:AddTags"
]
resources = [
"arn:aws:sagemaker:${local.region}:${local.account_id}:*/*"
]
condition {
test = "Null"
variable = "sagemaker:TaggingAction"
values = ["false"]
}
}
statement {
sid = "SMStudioRestrictSharedSpacesWithoutOwners"
effect = "Allow"
actions = [
"sagemaker:CreateSpace",
"sagemaker:UpdateSpace",
"sagemaker:DeleteSpace"
]
resources = [
"arn:aws:sagemaker:${local.region}:${local.account_id}:space/${aws_sagemaker_domain.this.id}/*",
]
condition {
test = "Null"
variable = "sagemaker:OwnerUserProfileArn"
values = ["true"]
}
}
statement {
sid = "SMStudioRestrictSpacesToOwnerUserProfile"
effect = "Allow"
actions = [
"sagemaker:CreateSpace",
"sagemaker:UpdateSpace",
"sagemaker:DeleteSpace"
]
resources = [
"arn:aws:sagemaker:${local.region}:${local.account_id}:space/${aws_sagemaker_domain.this.id}/*"
]
condition {
test = "ArnLike"
variable = "sagemaker:OwnerUserProfileArn"
values = [
aws_sagemaker_user_profile.this.arn
]
}
condition {
test = "StringEquals"
variable = "sagemaker:SpaceSharingType"
values = ["Private", "Shared"]
}
}
statement {
sid = "SMStudioRestrictCreatePrivateSpaceAppsToOwnerUserProfile"
effect = "Allow"
actions = [
"sagemaker:CreateApp",
"sagemaker:DeleteApp"
]
resources = [
"arn:aws:sagemaker:${local.region}:${local.account_id}:app/${aws_sagemaker_domain.this.id}/*"
]
condition {
test = "ArnLike"
variable = "sagemaker:OwnerUserProfileArn"
values = [
aws_sagemaker_user_profile.this.arn
]
}
condition {
test = "StringEquals"
variable = "sagemaker:SpaceSharingType"
values = ["Private"]
}
}
statement {
sid = "AllowAppActionsForSharedSpaces"
effect = "Allow"
actions = [
"sagemaker:CreateApp",
"sagemaker:DeleteApp"
]
resources = [
"arn:aws:sagemaker:*:*:app/${aws_sagemaker_domain.this.id}/*/*/*"
]
condition {
test = "StringEquals"
variable = "sagemaker:SpaceSharingType"
values = ["Shared"]
}
}
}
resource "aws_iam_policy" "studio_basic" {
name = "${local.prefix}-studio-basic-policy"
description = "IAM policy for SageMaker Studio Basic Permissions"
policy = data.aws_iam_policy_document.studio_basic.json
}
resource "aws_iam_role_policy_attachment" "studio_basic" {
role = aws_iam_role.studio.name
policy_arn = aws_iam_policy.studio_basic.arn
}
# For Slurm Orchestrator
# https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-studio-setup-slurm.html
data "aws_iam_policy_document" "studio_hyperpod" {
statement {
effect = "Allow"
actions = [
"ssm:StartSession",
"ssm:TerminateSession"
]
resources = [
"*"
]
}
statement {
effect = "Allow"
actions = [
"sagemaker:CreateCluster",
"sagemaker:ListClusters"
]
resources = [
"*"
]
}
statement {
effect = "Allow"
actions = [
"cloudwatch:PutMetricData",
"cloudwatch:GetMetricData"
]
resources = [
"*"
]
}
statement {
effect = "Allow"
actions = [
"sagemaker:DescribeCluster",
"sagemaker:DescribeClusterNode",
"sagemaker:ListClusterNodes",
"sagemaker:UpdateCluster",
"sagemaker:UpdateClusterSoftware"
]
resources = [
"arn:aws:sagemaker:${local.region}:${local.account_id}:cluster/*"
]
}
}
resource "aws_iam_policy" "studio_hyperpod" {
name = "${local.prefix}-studio-hyperpod-policy"
description = "IAM policy for SageMaker Studio Hyperpod Administration Permissions"
policy = data.aws_iam_policy_document.studio_hyperpod.json
}
resource "aws_iam_role_policy_attachment" "studio_hyperpod" {
role = aws_iam_role.studio.name
policy_arn = aws_iam_policy.studio_hyperpod.arn
}
オーケストレータによる権限の違いや基本的な SageMaker Studio の実行ロールは以下をご覧ください。
ドメイン/ユーザープロファイルの作成
SageMaker Studio を利用するため、ドメイン/ユーザープロファイルの作成を行います。
ユーザープロファイルに先ほど作成したロールが割り当てられる作成します。
###################################################
# SageMaker Domain
###################################################
resource "aws_sagemaker_domain" "this" {
domain_name = "${local.prefix}-hyperpod"
auth_mode = "IAM"
vpc_id = module.vpc.vpc_id
subnet_ids = module.vpc.private_subnets
default_user_settings {
execution_role = aws_iam_role.studio.arn
security_groups = [
aws_security_group.studio.id
]
}
default_space_settings {
execution_role = aws_iam_role.studio.arn
security_groups = [
aws_security_group.studio.id
]
}
}
resource "aws_sagemaker_user_profile" "this" {
domain_id = aws_sagemaker_domain.this.id
user_profile_name = "${local.prefix}-hyperpod-admin"
}
作成したユーザープロファイルから SageMaker Studio を起動します。
Compute から HyperPod clusters をクリックすると、現在利用している HyperPod クラスターの情報が表示されていますね。
Pytorch DDP
Pytorch を使った DDP を worker group に実行します。元ネタは SageMaker HyperPod のワークショップから持ってきました。
CloudShell から Controller Group にログイン後、 /fsx 配下で以下を実行しました。
# SSM Session Manager Plugin のインストール
sudo yum install -y https://s3.amazonaws.com/session-manager-downloads/plugin/latest/linux_64bit/session-manager-plugin.rpm
# Easy SSH のインストール
curl -O https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh
chmod +x easy-ssh.sh
# controller-group にログイン
./easy-ssh.sh -c controller-group tkkn-cluster
無事、ログインできていますね。
[cloudshell-user@ip-10-134-76-188 ~]$ # SSM Session Manager Plugin のインストール
[cloudshell-user@ip-10-134-76-188 ~]$ sudo yum install -y https://s3.amazonaws.com/session-manager-downloads/plugin/latest/linux_64bit/session-manager-plugin.rpm
# Easy SSH のインストール
curl -O https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh
chmod +x easy-ssh.sh
./easy-ssh.sh -c login-group tkkn-clusterLast metadata expiration check: 0:03:32 ago on Sat 28 Dec 2024 08:47:27 AM UTC.
session-manager-plugin.rpm 3.6 MB/s | 2.8 MB 00:00
Dependencies resolved.
==================================================================================================================================================================================================================================================================
Package Architecture Version Repository Size
==================================================================================================================================================================================================================================================================
Upgrading:
session-manager-plugin x86_64 1.2.694.0-1 @commandline 2.8 M
Transaction Summary
==================================================================================================================================================================================================================================================================
Upgrade 1 Package
Total size: 2.8 M
Downloading Packages:
Running transaction check
Transaction check succeeded.
Running transaction test
Transaction test succeeded.
Running transaction
Preparing : 1/1
Upgrading : session-manager-plugin-1.2.694.0-1.x86_64 1/2
Cleanup : session-manager-plugin-1.2.677.0-1.x86_64 2/2
Running scriptlet: session-manager-plugin-1.2.677.0-1.x86_64 2/2
Running scriptlet: session-manager-plugin-1.2.694.0-1.x86_64 2/2
Running scriptlet: session-manager-plugin-1.2.677.0-1.x86_64 2/2
Verifying : session-manager-plugin-1.2.694.0-1.x86_64 1/2
Verifying : session-manager-plugin-1.2.677.0-1.x86_64 2/2
Upgraded:
session-manager-plugin-1.2.694.0-1.x86_64
Complete!
[cloudshell-user@ip-10-134-76-188 ~]$ # Easy SSH のインストール
[cloudshell-user@ip-10-134-76-188 ~]$ curl -O https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh
% Total % Received % Xferd Average Speed Time Time Time Current
Dload Upload Total Spent Left Speed
100 5594 100 5594 0 0 21222 0 --:--:-- --:--:-- --:--:-- 21189
[cloudshell-user@ip-10-134-76-188 ~]$ chmod +x easy-ssh.sh
[cloudshell-user@ip-10-134-76-188 ~]$ ./easy-ssh.sh -c controller-group tkkn-cluster
=================================================
==== 🚀 HyperPod Cluster Easy SSH Script! 🚀 ====
=================================================
Cluster id: l2azofnyxblv
Instance id: i-014b35e0f513dd4fe
Node Group: controller-group
grep: /home/cloudshell-user/.ssh/config: No such file or directory
Would you like to add tkkn-cluster to ~/.ssh/config (yes/no)?
> yes
❌ skipping adding ml-cluster to ~/.ssh/config:
cat: /home/cloudshell-user/.ssh/id_rsa.pub: No such file or directory
2. Detected SSH public key ~/.ssh/id_rsa.pub on the cluster. Skipping adding...
Now you can run:
$ ssh tkkn-cluster
Starting session with SessionId: cm-takakuni.shinnosuke-eii2hqfhx3pnfnxo9cabfslhcu
#
Anaconda を使って DDP を実行します。
git clone https://github.com/aws-samples/awsome-distributed-training.git
cd awsome-distributed-training/3.test_cases/16.pytorch-cpu-ddp/slurm
bash 0.create-conda-env.sh
sbatch 1.conda-train.sbatch
squeue の実行結果から、投入したジョブが JOBID 10 で実行中であることがわかります。
root@ip-10-0-1-66:/usr/bin# squeue
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
10 dev cpu-ddp- root R 49:30 2 ip-10-0-1-[34,192]
SageMaker Studio に戻ります。
幾つがタブがあり Tasks をクリックすると、先ほど私が投入したタスクが表示されていますね。また、 Actions からジョブを操作することもできそうです。
続いてメトリクスです。単位が 1hr からなので粒度は大きいですが、簡易的にハードウェアメトリクスを監視するのに役立ちそうです。
もっと粒度を細かくみたい場合は Managed Grafana や CloudWatch のカスタムダッシュボードの出番でしょうか。
CPU/メモリ利用率に加えて、利用している/アイドル状態の CPU 数が見られる部分は嬉しいですね。
Settings では HuperPod クラスターに所属するインスタンスやインスタンスグループの詳細が閲覧できるような作りになっていました。
最後に Details です。ここは HyperPod のタグやログのリダイレクトリンクが記載されていました。（タグを付与していないため、わかりづらくてすみません）
まとめ
以上、「HyperPod in Studio のダッシュボードを眺めてみた」でした。
SageMaker Studio に統合することで、ユーザの権限管理を SageMaker Studio にまとめられる点は非常に魅力的なのではないでしょうか。
また 1hr の幅ですが、ハードウェアメトリクスのダッシュボードがみれる部分や、アイドル状態のリソース監視なども Slurm のマネジメントコンソールだと、ネイティブでサポートしていないため嬉しいですね。
このブログがどなたかの参考になれば幸いです。
AWS 事業本部コンサルティング部のたかくに（@takakuni_）でした！
